/*
    ref link:
        https://blog.csdn.net/u014338577/article/details/108060821

*/

#include <stdio.h>
#include <stdlib.h>

#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libswscale/swscale.h>
#include <libavdevice/avdevice.h>

static AVFormatContext * open_device(int *w, int *h, enum AVPixelFormat *pix_fmt) {
    int ret = 0;
    AVFormatContext *pfmt_ctx = NULL;
    AVInputFormat *infmt = NULL;
    AVCodecContext *pcodec_ctx = NULL;
    int video_idx = -1;
    
    avdevice_register_all();
    if ((pfmt_ctx = avformat_alloc_context()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avformat_alloc_context ERROR %s \n", av_err2str(AVERROR(ENOMEM)));
        return NULL;
    }

    if ((infmt = av_find_input_format("video4linux2")) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_find_input_format ERROR %s \n", av_err2str(AVERROR(EINVAL)));
        return NULL;
    }

    if((ret = avformat_open_input(&pfmt_ctx, "/dev/video0", infmt, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_open_input ERROR %s\n", av_err2str(AVERROR(EINVAL)));
        return NULL;
    }

    if ((ret = avformat_find_stream_info(pfmt_ctx, NULL)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avformat_find_stream_info ERROR %s\n", av_err2str(AVERROR(EINVAL)));
        return NULL;
    }

    for (int i=0; i<pfmt_ctx->nb_streams; i++) {
        if (pfmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_idx = i;
            *w = pfmt_ctx->streams[i]->codecpar->width;
            *h = pfmt_ctx->streams[i]->codecpar->height;
            break;
        }
    }

    if ((pcodec_ctx = avcodec_alloc_context3(NULL)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_alloc_context3 ERROR \n");
        return NULL;
    }

    if ((ret = avcodec_parameters_to_context(pcodec_ctx, pfmt_ctx->streams[video_idx]->codecpar)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_parameters_to_context ERROR \n");
        return NULL;
    }

    *pix_fmt = pcodec_ctx->pix_fmt;
    
    avcodec_free_context(&pcodec_ctx);

    av_dump_format(pfmt_ctx, 0, "/dev/video0", 0);
    return pfmt_ctx;
}

static AVCodecContext *open_encodec(int w, int h) {
    AVCodec *pcodec = NULL;
    AVCodecContext *pcodec_ctx = NULL;
    AVDictionary *pdict = NULL;
    int ret = 0;
    if ((pcodec = avcodec_find_encoder(AV_CODEC_ID_H264)) == NULL) { 
    //if ((pcodec = avcodec_find_encoder_by_name("libx264")) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_find_encoder_by_name ERROR \n");
        return NULL;
    }

    if ((pcodec_ctx = avcodec_alloc_context3(pcodec)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_alloc_context3 ERROR \n");
        return NULL;
    }

    //av_dict_set(&pdict, );
    pcodec_ctx->width = w;
    pcodec_ctx->height = h;
    pcodec_ctx->keyint_min = 50;
    pcodec_ctx->level = 20;
    pcodec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
    pcodec_ctx->time_base = (AVRational){1, 25};
    pcodec_ctx->framerate = (AVRational){25, 1};

    if ((ret = avcodec_open2(pcodec_ctx, pcodec, &pdict)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_open2 ERROR \n");
        return NULL;
    }

    return pcodec_ctx;
}

static int encode(AVCodecContext *pcodec_ctx, AVFrame *pfrm, AVPacket *ppkt, FILE *fp) {
    int ret = 0;
    if ((ret = avcodec_send_frame(pcodec_ctx, pfrm)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "avcodec_send_frame ERROR %s \n", av_err2str(ret));
        return ret;
    }

    while (ret >= 0) {
        ret = avcodec_receive_packet(pcodec_ctx, ppkt);
        if(ret == AVERROR_EOF || ret == AVERROR(EAGAIN)) {
            break;
        } else if (ret < 0) {
            av_log(NULL, AV_LOG_ERROR, "avcodec_receive_packet ERROR %s \n", av_err2str(ret));
            return ret;
        }
        fwrite(ppkt->data, 1, ppkt->size, fp);
        fflush(fp);
    }
    return 0;
}

int main(int argc, char **argv) {
    
    if (argc != 4) {
        fprintf(stdout, "Usage: ./camera_h264 ./out.h264 ./yuyv.yuv ./yuv420p.yuv \n");
        fprintf(stderr, "You must input h264 filepath and yuv filepath ! \n");
        return -1;
    }
    av_log_set_level(AV_LOG_DEBUG);

    AVFormatContext *pfmt_ctx = NULL;
    AVCodecContext *pcodec_ctx = NULL;
    AVPacket *ppkt = NULL;
    AVPacket *ph264pkt = NULL;
    AVFrame *pfrm = NULL; // forh264 coding, format yuv420p
    struct SwsContext *psws_ctx = NULL;
    uint8_t *data[4] = {NULL}; // buffer same size with camera input, camera input -> data -> pfrm -> x264 encoding
    int linesize[4] = {0}; // 

    int in_w = 0; //camera input width
    int in_h = 0; //camera input height
    enum AVPixelFormat in_pix_fmt; //camera input format yuv422:yuyv

    int ret = 0;
    FILE *yuyv_fp = NULL;
    FILE *yuv420p_fp = NULL;
    FILE *h264_fp = NULL;
    int count = 0;

    char *h264_file = argv[1];
    char *yuyv_file = argv[2];
    char *yuv420p_file = argv[3];

    fprintf(stdout, "INFO: Output h264 file:%s, Output yuyv file:%s, Output yuv420p file %s \n", h264_file, yuyv_file, yuv420p_file);

    int width = 352; //wanted width
    int height = 288; //wanted height

    if ((pfmt_ctx = open_device(&in_w, &in_h, &in_pix_fmt)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "open_device ERROR \n");
        goto failed;
    }

    if ((pcodec_ctx = open_encodec(width, height)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "open_encodec ERROR \n");
        goto failed;
    }
    
    if ((ppkt = av_packet_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_packet_alloc ERROR \n");
        goto failed;
    }

    if ((ph264pkt = av_packet_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_packet_alloc ERROR \n");
        goto failed;
    }

    if ((ret = av_image_alloc(data, linesize, in_w, in_h, in_pix_fmt, 1)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "av_image_alloc ERROR \n");
        goto failed;
    }

    if ((pfrm = av_frame_alloc()) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "av_frame_alloc ERROR \n");
        goto failed;
    }

    if ((ret = av_image_alloc(pfrm->data, pfrm->linesize, width, height, AV_PIX_FMT_YUV420P, 1)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "av_image_alloc ERROR \n");
        goto failed;
    }
    pfrm->width = width;
    pfrm->height = height;
    pfrm->format = AV_PIX_FMT_YUV420P;

    if ((psws_ctx = sws_getContext(in_w, in_h, in_pix_fmt, width, height, AV_PIX_FMT_YUV420P, SWS_FAST_BILINEAR, NULL, NULL, NULL)) == NULL) {
        av_log(NULL, AV_LOG_ERROR, "sws_getContext ERROR \n");
        goto failed;
    }
    
    if ((yuyv_fp = fopen(yuyv_file, "wb+")) == NULL) {
        fprintf(stderr, "fopen open file %s ERROR \n", yuyv_file);
    }

    if ((yuv420p_fp = fopen(yuv420p_file, "wb+")) == NULL) {
        fprintf(stderr, "fopen open file %s ERROR \n", yuv420p_file);
    }

    if ((h264_fp = fopen(h264_file, "wb+")) == NULL) {
        fprintf(stderr, "fopen open file %s ERROR \n", h264_file);
    }

    while ((ret = av_read_frame(pfmt_ctx, ppkt)) >= 0) {
        // write yuyv
        fwrite(ppkt->data, 1, ppkt->size, yuyv_fp);
        fflush(yuyv_fp);
        count++;
        av_log(NULL, AV_LOG_INFO, "Current frame is %d \n", count);

        if ((ret = av_image_fill_arrays(data, linesize, ppkt->data, in_pix_fmt, in_w, in_h, 1)) < 0) {
            av_log(NULL, AV_LOG_ERROR, "av_image_fill_arrays ERROR \n");
            goto failed;
        }

        // yuyv -> yuv420p, size may also be resized
        sws_scale(psws_ctx, (const uint8_t *const *)data, linesize, 0, in_h, pfrm->data, pfrm->linesize);

        fwrite(pfrm->data[0], 1, width * height, yuv420p_fp);
        fwrite(pfrm->data[1], 1, width * height/4, yuv420p_fp);
        fwrite(pfrm->data[2], 1, width * height/4, yuv420p_fp);
        fflush(yuv420p_fp);

        pfrm->pts = count;

        if ((ret = encode(pcodec_ctx, pfrm, ph264pkt, h264_fp)) < 0) {
            av_log(NULL, AV_LOG_ERROR, "encode ERROR %s", av_err2str(ret));
            goto failed;
        }   

        av_packet_unref(ppkt);
        if (count >= 500) {
            break;
        }
    }

    if ((ret = encode(pcodec_ctx, NULL, ph264pkt, h264_fp)) < 0) {
        av_log(NULL, AV_LOG_ERROR, "encode ERROR %s", av_err2str(ret));
        goto failed;
    }  
    
failed:
    if (pfmt_ctx != NULL) {
        //avformat_close_input();
        avformat_free_context(pfmt_ctx);
    }

    if (pcodec_ctx != NULL) {
        avcodec_close(pcodec_ctx);
    }

    if (ppkt != NULL) {
        av_packet_free(&ppkt);
    }

    if (ph264pkt != NULL) {
        av_packet_free(&ph264pkt);
    }

    /*
    if (data[0] != NULL) {
        av_freep(&data[0]);
    }
    */

    if (pfrm != NULL) {
        av_frame_free(&pfrm);
    }

    if (yuyv_fp != NULL) {
        fclose(yuyv_fp);
    }

    if (yuv420p_fp != NULL) {
        fclose(yuv420p_fp);
    }

    if (h264_fp != NULL) {
        fclose(h264_fp);
    }
    return 0;
}